import pandas as pd

df = pd.read_excel('./全部A股.xlsx')

print(df.columns)

lvjm = df[df['证券简称']=='立讯精密']
sqhy = df[df['证券简称']=='三七互娱']


lvjm_inds = lvjm['所属申万行业名称(2021)\n[交易日期] 最新收盘日\n[行业级别] 二级行业'].values[0]
sqhy_inds = sqhy['所属申万行业名称(2021)\n[交易日期] 最新收盘日\n[行业级别] 二级行业'].values[0]

print(lvjm_inds)
print(sqhy_inds)

# cond1: 立讯精密同行业公司
cond1 = df[df['所属申万行业名称(2021)\n[交易日期] 最新收盘日\n[行业级别] 二级行业'].isin([lvjm_inds])]
print('立讯精密同行业公司数量：', cond1.shape[0])
cond1.insert(len(cond1.columns), '性质', '立讯精密同行业')

# cond2: 三七互娱同行业公司
cond2 = df[df['所属申万行业名称(2021)\n[交易日期] 最新收盘日\n[行业级别] 二级行业'].isin([sqhy_inds])]
print('三七互娱同行业公司数量：', cond2.shape[0])
cond2.insert(len(cond2.columns), '性质', '三七互娱同行业')

# cond3: 申万一级分类为石油石化、有色金属、煤炭、通信的公司 且为国央企
cond3 = df[df['所属申万行业名称(2021)\n[交易日期] 最新收盘日\n[行业级别] '].isin(['石油石化', '有色金属', '煤炭', '通信'])]
cond3 = cond3[cond3['企业所有制性质\n[交易日期] 最新收盘日'].isin(['中央国有企业', '地方国有企业'])]
print('申万一级分类为石油石化、有色金属、煤炭、通信的公司数量：', cond3.shape[0])
cond3.insert(len(cond3.columns), '性质', '国央企&能源通信类')

# cond4: 公司注册地为南京市
cond4 = df[df['城市']=='南京市']
print('注册地为南京市的公司数量: ', cond4.shape[0])
cond4.insert(len(cond4.columns), '性质', '注册地在南京')

# cond5: 额外增加的
extra = [
    '002432.SZ',    # 九安医疗
    '601899.SH',    # 紫金矿业
    '601808.SH',    # 中海油服
    '600188.SH',    # 兖矿能源
    '601669.SH',    # 中国电建
    '601800.SH',    # 中国交建
    '601390.SH',    # 中国中铁
    '601186.SH',    # 中国铁建
    '601155.SH',    # 新城控股
    '600900.SH',    # 长江电力
    '601688.SH',    # 华泰证券
    '600837.SH',    # 海通证券
    ]
cond5 = df[df['证券代码'].isin(extra)]
print('大模型自选的公司数量有: ', cond5.shape[0])
cond5.insert(len(cond5.columns), '性质', '大模型自选')


ret = pd.concat([cond1, cond2, cond3, cond4, cond5], axis=0)

# 对ret进行聚合，实现group by 证券代码和证券简称，concat 性质
# 首先创建一个临时DataFrame来存储聚合结果
grouped = ret.groupby(['证券代码', '证券简称'])['性质'].apply(lambda x: '、'.join(x.unique())).reset_index()


print('最终结果数量：', grouped.shape[0])
print('性质聚合示例：')
print(grouped.head())

# 将聚合结果写入原始Excel文件的新sheet中
try:
    # 尝试读取Excel文件，检查是否存在'筛选结果'表
    with pd.ExcelFile('./全部A股.xlsx') as xls:
        if '筛选结果' in xls.sheet_names:
            # 如果表已存在，使用mode='a'并覆盖该表
            with pd.ExcelWriter('./全部A股.xlsx', mode='a', engine='openpyxl', if_sheet_exists='replace') as writer:
                grouped.to_excel(writer, sheet_name='筛选结果', index=False)
        else:
            # 如果表不存在，直接添加
            with pd.ExcelWriter('./全部A股.xlsx', mode='a', engine='openpyxl') as writer:
                grouped.to_excel(writer, sheet_name='筛选结果', index=False)
    print('聚合结果已成功写入到全部A股.xlsx的"筛选结果"sheet中')
except Exception as e:
    print(f'写入Excel时出错: {e}')


# 将证券代码写入aim_code.txt文件
grouped['证券代码'].apply(lambda x: x.split('.')[0]).to_csv('aim_code.txt', index=False, header=False)
print('证券代码已成功写入aim_code.txt文件')
